#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
遗传算法优化的PBS调度解决方案
Genetic Algorithm Enhanced PBS Optimization
"""

import pandas as pd
import numpy as np
import random
from typing import List, Tuple, Dict
from dataclasses import dataclass
import warnings
warnings.filterwarnings('ignore')

@dataclass
class Individual:
    """遗传算法个体"""
    chromosome: List[int]  # 车辆序列
    fitness: float = 0.0
    scores: Dict = None
    
class GeneticPBSOptimizer:
    """遗传算法优化的PBS调度器"""
    
    def __init__(self, population_size=100, generations=200, mutation_rate=0.1, crossover_rate=0.8):
        self.population_size = population_size
        self.generations = generations
        self.mutation_rate = mutation_rate
        self.crossover_rate = crossover_rate
        
        # 问题数据
        self.car_info = {}
        self.data = None
        
        # 遗传算法参数
        self.elite_size = int(population_size * 0.1)  # 精英保留
        self.tournament_size = 5  # 锦标赛选择
        
    def load_data(self):
        """加载数据"""
        self.data = pd.read_excel('附件1.xlsx')
        self.data.columns = ['sequence', 'config', 'color', 'drive']
        
        # 直接使用原始值，不映射
        for _, row in self.data.iterrows():
            self.car_info[row['sequence']] = {
                'config': row['config'],
                'color': row['color'],
                'drive': row['drive'],
                'original_color': row['color'],
                'original_drive': row['drive']
            }
        
        print(f"遗传算法：加载{len(self.data)}辆车数据")
        print(f"颜色分布: {dict(self.data['color'].value_counts())}")
        print(f"驱动分布: {dict(self.data['drive'].value_counts())}")
    
    def calculate_fitness(self, chromosome: List[int]) -> Tuple[float, Dict]:
        """计算适应度值"""
        # 创建调度矩阵
        schedule = self.create_schedule_matrix(chromosome)
        
        # 计算各目标得分
        scores = self.calculate_all_scores(chromosome, schedule)
        
        # 加权总分作为适应度
        fitness = scores['total_score']
        
        return fitness, scores
    
    def create_schedule_matrix(self, sequence: List[int]) -> np.ndarray:
        """创建调度矩阵"""
        n_cars = len(sequence)
        n_cols = max(100, (n_cars + 9) // 10)
        
        schedule = np.zeros((10, n_cols), dtype=int)
        
        car_index = 0
        for col in range(n_cols):
            for row in range(10):
                if car_index < n_cars:
                    schedule[row, col] = sequence[car_index]
                    car_index += 1
                else:
                    break
            if car_index >= n_cars:
                break
        
        return schedule
    
    def calculate_all_scores(self, sequence: List[int], schedule: np.ndarray) -> Dict:
        """计算所有目标得分"""
        # 获取序列信息
        color_sequence = [self.car_info[car]['original_color'] for car in sequence]
        drive_sequence = [self.car_info[car]['original_drive'] for car in sequence]
        
        # 计算各目标得分
        obj1_score, _ = self.calculate_color_switching_score(color_sequence)
        obj2_score, _ = self.calculate_drive_ratio_score(drive_sequence)
        obj3_score, _ = self.calculate_pbs_utilization_score(schedule)
        obj4_score, _, _ = self.calculate_time_score(schedule)
        
        # 加权总分
        weights = {'obj1': 0.4, 'obj2': 0.3, 'obj3': 0.2, 'obj4': 0.1}
        total_score = (obj1_score * weights['obj1'] + 
                      obj2_score * weights['obj2'] + 
                      obj3_score * weights['obj3'] + 
                      obj4_score * weights['obj4'])
        
        return {
            'color_switching': obj1_score,
            'drive_ratio': obj2_score,
            'pbs_utilization': obj3_score,
            'total_time': obj4_score,
            'total_score': total_score
        }
    
    def calculate_color_switching_score(self, color_sequence: List[str]) -> Tuple[float, int]:
        """计算颜色切换得分"""
        switches = 0
        for i in range(1, len(color_sequence)):
            if color_sequence[i] != color_sequence[i-1]:
                switches += 1
        
        max_switches = len(color_sequence) - 1
        if max_switches > 0:
            score = max(0, 100 - (switches / max_switches) * 100)
        else:
            score = 100
            
        return score, switches
    
    def calculate_drive_ratio_score(self, drive_sequence: List[str]) -> Tuple[float, float]:
        """计算驱动比例得分"""
        total_deviation = 0
        window_size = 4
        
        for i in range(0, len(drive_sequence), window_size):
            window = drive_sequence[i:i+window_size]
            if len(window) < 2:
                continue
                
            two_count = window.count('两驱')
            four_count = window.count('四驱')
            
            ideal_two = len(window) / 2
            ideal_four = len(window) / 2
            
            deviation = abs(two_count - ideal_two) + abs(four_count - ideal_four)
            total_deviation += deviation
            
        max_possible_deviation = len(drive_sequence)
        if max_possible_deviation > 0:
            score = max(0, 100 - (total_deviation / max_possible_deviation) * 100)
        else:
            score = 100
            
        return score, total_deviation
    
    def calculate_pbs_utilization_score(self, schedule: np.ndarray) -> Tuple[float, float]:
        """计算PBS利用率得分"""
        used_positions = np.sum(schedule > 0)
        total_positions = schedule.shape[0] * schedule.shape[1]
        utilization = used_positions / total_positions
        score = utilization * 100
        return score, utilization
    
    def calculate_time_score(self, schedule: np.ndarray) -> Tuple[float, float, float]:
        """计算时间得分"""
        max_col_time = 0
        for col in range(schedule.shape[1]):
            col_cars = schedule[:, col]
            col_cars = col_cars[col_cars > 0]
            if len(col_cars) > 0:
                col_time = len(col_cars) * 9 + 12
                max_col_time = max(max_col_time, col_time)
        
        total_cars = np.sum(schedule > 0)
        min_time = 9 * total_cars + 72
        
        time_deviation = max(0, max_col_time - min_time)
        time_penalty = time_deviation * 0.1
        score = max(0, 100 - time_penalty)
        
        return score, max_col_time, min_time
    
    def initialize_population(self) -> List[Individual]:
        """初始化种群"""
        n_cars = len(self.data)
        base_sequence = list(range(1, n_cars + 1))
        population = []
        
        # 创建多种初始化策略
        strategies = [
            lambda: base_sequence.copy(),  # 原始顺序
            lambda: self.heuristic_color_first(),  # 颜色优先
            lambda: self.heuristic_drive_balanced(),  # 驱动平衡
            lambda: random.sample(base_sequence, n_cars),  # 随机
        ]
        
        for i in range(self.population_size):
            strategy = strategies[i % len(strategies)]
            chromosome = strategy()
            
            fitness, scores = self.calculate_fitness(chromosome)
            population.append(Individual(chromosome, fitness, scores))
        
        return population
    
    def heuristic_color_first(self) -> List[int]:
        """颜色优先启发式"""
        color_groups = {'电动': [], '燃动': []}
        for car_id, info in self.car_info.items():
            color_groups[info['original_color']].append(car_id)
        
        # 合并分组
        sequence = color_groups['电动'] + color_groups['燃动']
        return sequence
    
    def heuristic_drive_balanced(self) -> List[int]:
        """驱动平衡启发式"""
        two_drive = [car_id for car_id, info in self.car_info.items() 
                    if info['original_drive'] == '两驱']
        four_drive = [car_id for car_id, info in self.car_info.items() 
                     if info['original_drive'] == '四驱']
        
        # 交错排列
        sequence = []
        two_ptr, four_ptr = 0, 0
        
        while two_ptr < len(two_drive) or four_ptr < len(four_drive):
            if two_ptr < len(two_drive):
                sequence.append(two_drive[two_ptr])
                two_ptr += 1
            if four_ptr < len(four_drive):
                sequence.append(four_drive[four_ptr])
                four_ptr += 1
        
        return sequence
    
    def tournament_selection(self, population: List[Individual]) -> Individual:
        """锦标赛选择"""
        tournament = random.sample(population, self.tournament_size)
        return max(tournament, key=lambda x: x.fitness)
    
    def crossover(self, parent1: List[int], parent2: List[int]) -> Tuple[List[int], List[int]]:
        """顺序交叉(OX)"""
        n = len(parent1)
        
        # 随机选择交叉点
        start, end = sorted(random.sample(range(n), 2))
        
        # 创建子代
        child1 = [-1] * n
        child2 = [-1] * n
        
        # 复制交叉段
        child1[start:end+1] = parent1[start:end+1]
        child2[start:end+1] = parent2[start:end+1]
        
        # 填充剩余位置
        def fill_child(child, parent):
            ptr = (end + 1) % n
            for gene in parent:
                if gene not in child:
                    child[ptr] = gene
                    ptr = (ptr + 1) % n
            return child
        
        child1 = fill_child(child1, parent2)
        child2 = fill_child(child2, parent1)
        
        return child1, child2
    
    def mutate(self, chromosome: List[int]) -> List[int]:
        """交换变异"""
        if random.random() < self.mutation_rate:
            # 随机选择两个位置交换
            i, j = random.sample(range(len(chromosome)), 2)
            chromosome[i], chromosome[j] = chromosome[j], chromosome[i]
        return chromosome
    
    def evolve(self) -> Tuple[List[int], Dict, List[float]]:
        """遗传算法主循环"""
        print("=== 遗传算法优化开始 ===")
        
        # 初始化种群
        population = self.initialize_population()
        
        # 记录最佳适应度历史
        fitness_history = []
        
        for generation in range(self.generations):
            # 评估适应度
            for individual in population:
                if individual.fitness == 0:  # 未评估
                    fitness, scores = self.calculate_fitness(individual.chromosome)
                    individual.fitness = fitness
                    individual.scores = scores
            
            # 排序
            population.sort(key=lambda x: x.fitness, reverse=True)
            
            # 记录最佳
            best_fitness = population[0].fitness
            fitness_history.append(best_fitness)
            
            if generation % 20 == 0:
                print(f"第{generation}代: 最佳适应度={best_fitness:.2f}")
            
            # 创建新一代
            new_population = population[:self.elite_size]  # 保留精英
            
            # 生成后代
            while len(new_population) < self.population_size:
                parent1 = self.tournament_selection(population)
                parent2 = self.tournament_selection(population)
                
                if random.random() < self.crossover_rate:
                    child1_chr, child2_chr = self.crossover(parent1.chromosome, parent2.chromosome)
                else:
                    child1_chr, child2_chr = parent1.chromosome.copy(), parent2.chromosome.copy()
                
                # 变异
                child1_chr = self.mutate(child1_chr)
                child2_chr = self.mutate(child2_chr)
                
                # 计算适应度
                fitness1, scores1 = self.calculate_fitness(child1_chr)
                fitness2, scores2 = self.calculate_fitness(child2_chr)
                
                new_population.extend([
                    Individual(child1_chr, fitness1, scores1),
                    Individual(child2_chr, fitness2, scores2)
                ])
            
            population = new_population[:self.population_size]
        
        # 最终结果
        population.sort(key=lambda x: x.fitness, reverse=True)
        best_individual = population[0]
        
        print(f"=== 遗传算法优化完成 ===")
        print(f"最终适应度: {best_individual.fitness:.2f}")
        
        return best_individual.chromosome, best_individual.scores, fitness_history
    
    def generate_results(self, sequence: List[int], scores: Dict):
        """生成结果文件"""
        # 创建调度矩阵
        schedule = self.create_schedule_matrix(sequence)
        
        # 生成车辆序列信息
        sequence_data = []
        for car_id in sequence:
            if car_id > 0 and car_id in self.car_info:
                car_data = self.car_info[car_id]
                sequence_data.append({
                    '车辆顺序': car_id,
                    '车型': car_data['config'],
                    '动力': car_data['original_color'],
                    '驱动': car_data['original_drive']
                })
        
        # 保存结果
        schedule_df = pd.DataFrame(schedule)
        sequence_df = pd.DataFrame(sequence_data)
        
        schedule_df.to_excel('result11_genetic.xlsx', index=False, header=False)
        sequence_df.to_excel('result12_genetic.xlsx', index=False)
        
        # 保存得分
        score_df = pd.DataFrame([scores])
        score_df.to_excel('genetic_scores.xlsx', index=False)
        
        print("遗传算法结果已保存:")
        print("- result11_genetic.xlsx (调度矩阵)")
        print("- result12_genetic.xlsx (车辆序列)")
        print("- genetic_scores.xlsx (得分详情)")

def main():
    """主函数"""
    optimizer = GeneticPBSOptimizer(
        population_size=100,
        generations=100,
        mutation_rate=0.1,
        crossover_rate=0.8
    )
    
    optimizer.load_data()
    best_sequence, best_scores, history = optimizer.evolve()
    optimizer.generate_results(best_sequence, best_scores)
    
    return best_sequence, best_scores, history

if __name__ == "__main__":
    main()